import os
import asyncio
from typing import Optional, List
import httpx, requests
from mcp.server.fastmcp import FastMCP


def get_api_key() -> str:
    """Get API key from environment variables"""
    api_key = os.getenv("ZHIPUAI_API_KEY")
    if not api_key:
        raise ValueError("ZHIPUAI_API_KEY环境变量未设置，请检查客户端配置")
    return api_key


mcp = FastMCP("BigModel Media MCP Server")


@mcp.tool()
async def text_to_video(
    prompt: str,
    quality: str = "speed",
    with_audio: bool = False,
    size: str = "1920x1080",
    fps: int = 30,
    duration: int = 5,
    request_id: Optional[str] = None,
    user_id: Optional[str] = None,
) -> str:
    """
    文生视频 - 通过文本提示生成视频 (CogVideoX-3)

    Args:
        prompt: 文本描述要生成的视频内容
        quality: 输出模式，"quality"为质量优先，"speed"为速度优先
        with_audio: 是否生成AI音效
        size: 视频分辨率，如"1920x1080"
        fps: 帧率，可选30或60
        duration: 视频时长(秒)，支持5或10
        request_id: 用户端传参，需保证唯一性
        user_id: 终端用户的唯一ID

    Returns:
        task_id: 用于查询任务结果的ID
    """
    API_URL = "https://open.bigmodel.cn/api/paas/v4/videos/generations"

    api_key = get_api_key()
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
    }

    data = {
        "model": "cogvideox-3",
        "prompt": prompt,
        "quality": quality,
        "with_audio": with_audio,
        "size": size,
        "fps": fps,
        "duration": duration,
    }

    if request_id:
        data["request_id"] = request_id
    if user_id:
        data["user_id"] = user_id

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(API_URL, headers=headers, json=data)
            response.raise_for_status()

            result = response.json()
            return result["id"]
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def image_to_video(
    image_url: str,
    prompt: Optional[str] = None,
    quality: str = "speed",
    with_audio: bool = False,
    size: str = "1920x1080",
    fps: int = 30,
    duration: int = 5,
    request_id: Optional[str] = None,
    user_id: Optional[str] = None,
) -> str:
    """
    图生视频 - 基于图片生成视频 (CogVideoX-3)

    Args:
        image_url: 图片URL或Base64编码
        prompt: 可选文本描述
        quality: 输出模式，"quality"为质量优先，"speed"为速度优先
        with_audio: 是否生成AI音效
        size: 视频分辨率，如"1920x1080"
        fps: 帧率，可选30或60
        duration: 视频时长(秒)，支持5或10
        request_id: 用户端传参，需保证唯一性
        user_id: 终端用户的唯一ID

    Returns:
        task_id: 用于查询任务结果的ID
    """
    API_URL = "https://open.bigmodel.cn/api/paas/v4/videos/generations"

    api_key = get_api_key()
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
    }

    data = {
        "model": "cogvideox-3",
        "image_url": image_url,
        "quality": quality,
        "with_audio": with_audio,
        "size": size,
        "fps": fps,
        "duration": duration,
    }

    if prompt:
        data["prompt"] = prompt
    if request_id:
        data["request_id"] = request_id
    if user_id:
        data["user_id"] = user_id

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(API_URL, headers=headers, json=data)
            response.raise_for_status()

            result = response.json()
            return result["id"]
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def first_last_frame_video(
    first_frame_url: str,
    last_frame_url: str,
    prompt: Optional[str] = None,
    with_audio: bool = False,
    size: str = "1920x1080",
    fps: int = 30,
    duration: int = 5,
    request_id: Optional[str] = None,
    user_id: Optional[str] = None,
) -> str:
    """
    首尾帧视频 - 基于首尾帧图片生成视频 (CogVideoX-3)

    Args:
        first_frame_url: 首帧图片URL或Base64编码
        last_frame_url: 尾帧图片URL或Base64编码
        prompt: 可选文本描述
        with_audio: 是否生成AI音效
        size: 视频分辨率，如"1920x1080"
        fps: 帧率，可选30或60
        duration: 视频时长(秒)，支持5或10
        request_id: 用户端传参，需保证唯一性
        user_id: 终端用户的唯一ID

    Returns:
        task_id: 用于查询任务结果的ID
    """
    API_URL = "https://open.bigmodel.cn/api/paas/v4/videos/generations"

    api_key = get_api_key()
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
    }

    data = {
        "model": "cogvideox-3",
        "image_url": [first_frame_url, last_frame_url],
        "quality": "speed",  # 首尾帧模式只支持speed
        "with_audio": with_audio,
        "size": size,
        "fps": fps,
        "duration": duration,
    }

    if prompt:
        data["prompt"] = prompt
    if request_id:
        data["request_id"] = request_id
    if user_id:
        data["user_id"] = user_id

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(API_URL, headers=headers, json=data)
            response.raise_for_status()

            result = response.json()
            return result["id"]
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
def get_video_result(task_id: str) -> dict:
    """
    查询视频生成结果 (CogVideoX-3)

    Args:
        task_id: 任务ID

    Returns:
        dict: 包含视频URL和封面URL的字典

    Raises:
        ValueError: 当任务ID无效或任务失败时
        requests.RequestException: 网络请求失败时
    """
    API_URL = f"https://open.bigmodel.cn/api/paas/v4/async-result/{task_id}"

    api_key = get_api_key()
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
    }

    try:
        response = requests.get(API_URL, headers=headers)
        response.raise_for_status()

        result = response.json()
        if result["task_status"] == "SUCCESS":
            return {
                "video_url": result["video_result"][0]["url"],
                "cover_url": result["video_result"][0]["cover_image_url"],
            }
        elif result["task_status"] == "FAIL":
            raise ValueError("视频生成失败")
        else:
            return {"status": result["task_status"]}

    except requests.RequestException as e:
        raise requests.RequestException(f"请求失败: {str(e)}")


@mcp.tool()
def generate_image(
    prompt: str,
    model: str = "cogview-4-250304",
    quality: str = "standard",
    size: str = "1024x1024",
    user_id: Optional[str] = None,
) -> str:
    """
    图像生成 - 通过文本描述生成图像 (CogView-4)

    Args:
        prompt: 图像描述的文本，长度不超过2000字符
        model: 模型编码，支持 cogview-4-250304(最新)、cogview-4、cogview-3-flash
        quality: 生成质量，"hd"为高质量(耗时约20秒)，"standard"为标准质量(5-10秒)
        size: 图片尺寸，支持: 1024x1024,768x1344,864x1152,1344x768,1152x864,1440x720,720x1440
              或自定义尺寸(512-2048px之间，能被16整除)
        user_id: 终端用户的唯一ID(6-128字符)，用于安全审计

    Returns:
        str: 生成的图片URL(30天有效)

    Raises:
        ValueError: 参数验证失败时
        requests.RequestException: API请求失败时
    """
    # 验证size参数
    try:
        width, height = map(int, size.split("x"))
        if not (512 <= width <= 2048 and 512 <= height <= 2048):
            raise ValueError("图片尺寸必须在512-2048像素之间")
        if width % 16 != 0 or height % 16 != 0:
            raise ValueError("图片尺寸必须能被16整除")
        if width * height > 2**21:  # 2^21 = 2097152
            raise ValueError("图片总像素数不能超过2097152")
    except Exception as e:
        raise ValueError(f"无效的size参数: {str(e)}")

    API_URL = "https://open.bigmodel.cn/api/paas/v4/images/generations"

    api_key = get_api_key()
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
    }

    data = {
        "model": model,
        "prompt": prompt,
        "quality": quality,
        "size": size,
    }

    if user_id:
        if not (6 <= len(user_id) <= 128):
            raise ValueError("user_id长度必须在6-128字符之间")
        data["user_id"] = user_id

    try:
        response = requests.post(API_URL, headers=headers, json=data)
        response.raise_for_status()

        result = response.json()

        # 检查内容安全
        if "content_filter" in result:
            for item in result["content_filter"]:
                if item["level"] == 0:  # 最严重级别
                    raise ValueError("内容安全检测不通过: 包含违规内容")

        return result["data"][0]["url"]

    except requests.RequestException as e:
        raise requests.RequestException(f"API请求失败: {str(e)}")
    except KeyError as e:
        raise ValueError(f"无效的API响应: 缺少 {str(e)} 字段")
    except Exception as e:
        raise ValueError(f"处理API响应时出错: {str(e)}")


def main():
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
